import tensorflow as tf

'''
工具类：
获取张量形状、张量展平与逆展平、计算模型性能指标、获取数据集维度
'''

def get_static_shape(tensor):
    """
    获取指定张量的形状

    参数:
        tensor: tensor 对象

    返回:
        tuple[int or None] or None: 张量形状
        如果 tensor 维度不确定，则返回None
    """
    # 将输入转换为 TensorFlow 的张量对象
    tensor = tf.convert_to_tensor(tensor)
    # 获取张量形状
    shape = tensor.get_shape()
    # 检查张量的维度数，如果未知，则将形状设为 None 如果已知则返回整型元组
    if shape.ndims is None:
        shape = None
    else:
        shape = tuple((int(v) if v is not None else None)
                      for v in shape.as_list())
    return shape

def flatten_to_ndims(x, ndims):
    """
    将输入 `x` 展平到 `ndims` 个维度

    参数：
        x (tensor): 需要展平的张量
        ndims (int): 结果张量的最大维度数

    返回：
        (tf.tensor, tuple[int or None], tuple[int] 或 tf.tensor) 或 (tf.tensor, None, None):
        (展平后的张量，静态前形状，和前形状)，或 (原始张量, None, None)
    """
    # 将输入张量 x 转换为 tensorflow 张量对象
    x = tf.convert_to_tensor(x)
    # 检查 ndims 的值，确保其大于或等于1，否则抛出异常
    if ndims < 1:
        raise ValueError('`k` must be greater or equal to 1.')
    # 检查张量的形状信息，确保其已知，否则抛出异常
    if not x.get_shape():
        raise ValueError('`x` is required to have known number of dimensions.')
    # 获取张量的静态形状
    shape = get_static_shape(x)
    # 如果张量的维度数小于 ndims，则抛出异常
    if len(shape) < ndims:
        raise ValueError('`k` is {}, but `x` only has rank {}.'.format(ndims, len(shape)))
    # 如果维度数等于 ndims，返回原始张量和两个 None
    if len(shape) == ndims:
        return x, None, None
    # 处理一维张量的情况，返回重塑后的张量和相关形状信息
    if ndims == 1:
        static_shape = shape
        if None in shape:
            shape = tf.shape(x)
        return tf.reshape(x, [-1]), static_shape, shape
    # 处理多维情况
    else:
        front_shape, back_shape = shape[:-(ndims - 1)], shape[-(ndims - 1):]
        static_front_shape = front_shape
        static_back_shape = back_shape
        if None in front_shape or None in back_shape:
            dynamic_shape = tf.shape(x)
            if None in front_shape:
                front_shape = dynamic_shape[:-(ndims - 1)]
            if None in back_shape:
                back_shape = dynamic_shape[-(ndims - 1):]
        if isinstance(back_shape, tuple):
            x = tf.reshape(x, [-1] + list(back_shape))
        else:
            x = tf.reshape(x, tf.concat([[-1], back_shape], axis=0))
            x.set_shape(tf.TensorShape([None] + list(static_back_shape)))
        return x, static_front_shape, front_shape


def unflatten_from_ndims(x, static_front_shape, front_shape):
    """
    函数 `flatten` 的逆变换

    如果 `static_front_shape` 和 `front_shape` 都为 None，则 `x` 将不做任何改变地返回。

    参数：
        x (tensor): 需要逆展平的张量
        static_front_shape (tuple[int or None] 或 None): 静态形状
        front_shape (tuple[int] 或 tf.tensor 或 None): 展平后形状

    返回：
        x(tf.tensor): 逆展平后的 `x`
    """
    # 将输入张量转换为 TensorFlow 张量对象
    x = tf.convert_to_tensor(x)
    # 如果两个形状参数都为 None，直接返回原始张量
    if static_front_shape is None and front_shape is None:
        return x
    # 检查张量 x 是否有已知形状  如果没有，则抛出一个 ValueError 异常
    if not x.get_shape():
        raise ValueError('`x` is required to have known number of dimensions.')
    # 获取形状
    shape = get_static_shape(x)
    if len(shape) < 1:
        raise ValueError('`x` only has rank {}, required at least 1.'.
                         format(len(shape)))

    if not isinstance(front_shape, tf.Tensor):
        front_shape = tuple(front_shape)

    back_shape = shape[1:]
    static_back_shape = back_shape
    if None in back_shape:
        back_shape = tf.shape(x)[1:]
    if isinstance(front_shape, tuple) and isinstance(back_shape, tuple):
        x = tf.reshape(x, front_shape + back_shape)
    else:
        x = tf.reshape(x, tf.concat([front_shape, back_shape], axis=0))
        x.set_shape(tf.TensorShape(list(static_front_shape) + list(static_back_shape)))
    return x

def get_best_f1(score, label):
    '''
    计算给定评分和标签的最佳 F1 分数及其对应的阈值，用于评估模型在异常检测任务中的性能
    score: 一维数组，表示模型的评分
    label: 一维数组，表示标准标签（指示异常）
    返回
    '''
    assert score.shape == label.shape
    print('***computing best f1***')
    search_set = []
    tot_anomaly = 0
    for i in range(label.shape[0]):
        tot_anomaly += (label[i] > 0.5)
    flag = 0
    cur_anomaly_len = 0
    cur_min_anomaly_score = 1e5
    for i in range(label.shape[0]):
        if label[i] > 0.5:
            # here for an anomaly
            if flag == 1:
                cur_anomaly_len += 1
                cur_min_anomaly_score = score[i] if score[i] < cur_min_anomaly_score else cur_min_anomaly_score
            else:
                flag = 1
                cur_anomaly_len = 1
                cur_min_anomaly_score = score[i]
        else:
            # here for normal points
            if flag == 1:
                flag = 0
                search_set.append((cur_min_anomaly_score, cur_anomaly_len, True))
                search_set.append((score[i], 1, False))
            else:
                search_set.append((score[i], 1, False))
    if flag == 1:
        search_set.append((cur_min_anomaly_score, cur_anomaly_len, True))
    search_set.sort(key=lambda x: x[0])
    best_f1_res = - 1
    threshold = 1
    P = 0
    TP = 0
    best_P = 0
    best_TP = 0
    for i in range(len(search_set)):
        P += search_set[i][1]
        if search_set[i][2]:  # for an anomaly point
            TP += search_set[i][1]
        precision = TP / (P + 1e-5)
        recall = TP / (tot_anomaly + 1e-5)
        f1 = 2 * precision * recall / (precision + recall + 1e-5)
        if f1 > best_f1_res:
            best_f1_res = f1
            threshold = search_set[i][0]
            best_P = P
            best_TP = TP

    print('***  best_f1  ***: ', best_f1_res)
    print('*** threshold ***: ', threshold)

    return (best_f1_res,
            best_TP / (best_P + 1e-5),
            best_TP / (tot_anomaly + 1e-5),
            best_TP,
            score.shape[0] - best_P - tot_anomaly + best_TP,
            best_P - best_TP,
            tot_anomaly - best_TP), threshold

def get_data_dim(dataset):
    # 根据数据集名称返回其特征的维度数
    if str(dataset).startswith('machine'):
        return 38
    elif str(dataset).startswith('omi'):
        return 19
    elif dataset == 'PSM':
        return 25
    else:
        raise ValueError('unknown dataset '+str(dataset))

def get_data_name(dataset):
    if str(dataset).startswith('omi'):
        return 'ASD'
    elif str(dataset).startswith('machine'):
        return 'SMD'
    else:
        return dataset